# -*- coding: utf-8 -*-
import torch.nn as nn
import torch.nn.functional as F


class ResidualFeedForward(nn.Module):
    def __init__(self, d_model, hidden=2048, dropout=0.1):
        super(ResidualFeedForward, self).__init__()
        self.linear1 = nn.Linear(d_model, hidden, bias=False)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(hidden, d_model, bias=False)
        self.layer_normal = nn.LayerNorm(d_model)

    def forward(self, x):
        """

        :param x:       B, L, H
        :return:
        """
        # feed-forward
        residual = x
        x = self.dropout(F.relu(self.linear1(x)))
        x = self.linear2(x)

        # add & norm
        x = residual + x  # Add
        x = self.layer_normal(x)  # Norm
        return x
